[Perf] Reduce DFlash2 long-context verification decay - #587

Purpose

Reduce complete DFlash2 verification-round latency and long-context growth on rear GPUs 4–7, TP4/B1/q8. Targets are 32K ≤17 ms, 64K ≤18 ms, 128K <20 ms and 256K <22 ms, with no quality or acceptance regression. Targets remain unmet; all experimental defaults remain disabled. Service capacity remains 262144; the boundary-window test uses 261888 input plus 256 output tokens.

This continues #587, stacked on #586, without opening a duplicate PR. Integration freeze: 80545c010bbf6f5ed06458d992c189d75d0eff8f. Serving source remains 239d71c7100b3bce5526268be2cafb4cff8ba8f2; review head is 1994baff8f188a2c915e883395ae96e682db0577. Source/library hashes, dispatch and retained failures are in the worklog.

Implementation

• Opt-in six-head KV reuse, vector loads, page specialization and ordered QK/PV scheduling with MRV2 context descriptors. Preserve 80 logical splits, N32 updates, K16 compensation, probability residual products and FP32 state/logits. Shape/layout admission is independent of model name or weight quantization.
• Add scalar q1 shared-KV/lookup probes and compact page-map resource controls, with unchanged scalar FP32 operation and partition order.
• Add explicit native pre-cast FP32 reference diagnostics and complete small-Q coverage. Reject failed scheduling/arithmetic candidates; none enables a default.

Test Plan and Result

Visible-tile q8 versus its PV-reuse parent completes three paired independent starts, 180 requests and 90 exact pairs, including tokens, finish reason and acceptance. One cold and five measured requests per context/arm; no profiler or tensor dump. Original scalar q1 is retained in both arms:

These are medians of startup medians. The third 32K pair is 0.005 ms slower and remains included. Request-average quantiles, cold TTFT/prefill and incremental costs are separately reported in the worklog. Native output/full-workspace checks and the visible operator's own three sanitizers pass.

Compact scalar lookup plus a 4096-byte shared reservation passes 45 operator byte checks and all three sanitizers. Its first service A/B holds visible-tile q8 fixed, hits 384 scalar calls per rank and preserves all 12 request pairs: 261888 round 37.160→35.602 ms, pure decode 137.245→143.250 tokens/s, accepted drafts 4.02/round and emitted tokens 5.12/round unchanged. This is one startup; it cannot be ranked against a different earlier trajectory. A separate actual-input diagnostic retains 96 exact comparisons but samples only eight unique K/V pointer pairs per rank; its incorrect sixteen-pair final assertion and incomplete full-layer coverage remain recorded.

New full-round traces confirm q6 and q1 eager work at the capacity boundary. Every q2–q7 shape now passes 60 output/full-FP32-workspace checks against the actual frozen production library. The q6 16-layer working set improves 12.176→7.796 ms at 128K and 23.717→15.066 ms at 261888. This is an operator screen; small-Q sanitizers and service-route admission are still required.

The original selected natural-EOS campaign has 26/30 complete exact pairs, including all 12 structured/tool cases; four seed-2 code cases remain. These are paired non-regression observations, not new benchmark scores. Unchanged-control startup drift remains unresolved and is never an allowed numerical tolerance.

Scoped pre-commit passes. GPU commands use benchmark_sm70_grouped_attention_long, benchmark_sm70_scalar_attention, and benchmark_sm70_grouped_attention_precision; source/library identities and raw reports are retained. NCU counters remain unavailable. Static resources, CTA clocks and data-derived bank proxies are not achieved hardware utilization.

Remaining gates

Finish tail-route attribution and actual-input coverage, then paired service/repeated-start admission for the best tail combination. Continue reducing the major costs exposed by the complete trace. N64 variants, register-held softmax state, M8/N32 QK, operand rotation and the FP64-sum candidate fail their gates and remain closed. The user explicitly approved merging the current implementation on 2026-09-10. Default enablement remains subject to the outstanding quality and performance gates.

AI assistance: implementation, audit tooling and this report were assisted by Codex. The current main merge is explicitly user-authorized; numerical and performance claims retain the limits above.

Main integration

The user explicitly requested merging the current DFlash2 optimization work to main on 2026-09-10. Merge order: #586 → #587 → #589. The combined tree against main 0a0d4d6754f553ab00307eac807e9b5963b58b37 merges cleanly and passes CUDA_VISIBLE_DEVICES='' .venv/bin/python -m pytest -q tests/v1/worker/test_sm70_long_attention_graphs.py tests/kernels/core/test_sm70_dflash2_state_audit.py (26 passed, 1 skipped) and scoped pre-commit. This merge preserves explicit opt-in switches and does not assert the unfinished 7-ms attention or 22-ms round targets. It adds no default activation or new service admission.